"""
-*- coding: utf-8 -*-

@author: Du Changping
@time: 2021/12/10 15:59
@file name: create_vox_filev2
@software：PyCharm

Do not smash your computer!

this file is used to create vox from the existed phantom
"""

import os
import time
import argparse
import numpy as np
import SimpleITK as sitk
from tqdm import tqdm
import pandas as pd
import matplotlib.pyplot as plt


def path_setting(model_name_: str):
    """

    :param model_name_: model name to create the folder
    :return: current dir, the folder path to save the vox file
    """
    cur_dir_ = os.path.dirname(__file__)
    model_dir_path_ = os.path.join(cur_dir_, model_name_)
    if not os.path.exists(model_dir_path_):
        os.mkdir(model_dir_path_)
    return cur_dir_, model_dir_path_


def add_timestamp():
    time_stamp_ = time.strftime("%m%d-%H%M", time.localtime())
    return time_stamp_


def create_3d_vox_file(scale_: float, model_name_: str,
                       make_air: bool,
                       pixel_dist, time_stamp_="",
                       nii_file_path_=None):
    cur_dir_, model_save_dir_path_ = path_setting(model_name_)
    phantom_ = sitk.ReadImage(nii_file_path_)
    phantom_ = sitk.GetArrayFromImage(phantom_)
    dim_ = phantom_.shape
    assert len(dim_) > 2, "here we only support 3d object"

    nz_ = dim_[0]  # matrix dimension
    ny_ = dim_[2]
    nx_ = dim_[1]
    pixel_dist_list = pixel_dist.replace(" ", "").split(",")
    print("pixel distance is:", pixel_dist_list[0], pixel_dist_list[1], pixel_dist_list[2])
    dx_ = float(pixel_dist_list[0]) * scale_ * 0.1  # convert to cm
    dy_ = float(pixel_dist_list[1]) * scale_ * 0.1
    dz_ = float(pixel_dist_list[2]) * scale_ * 0.1

    phantom_ = phantom_.flatten(order="C")

    total_voxels_ = phantom_.size

    log_path_ = os.path.join(model_save_dir_path_, "{}_{}_{}.txt".format(model_name_, scale_, time_stamp_))
    print("nii phantom shape is {}".format(dim_))
    print("phantom size is {}. unit: cm".format([nx_ * dx_, ny_ * dy_, nz_ * dz_]))
    log_f_ = open(log_path_, "w")
    log_f_.write("nii dimension is {}\n".format(dim_))
    log_f_.write("pixel distance is [{}]\n".format(pixel_dist))
    log_f_.write("target object sizes is {}, {}, {}, (cm)\n".format(dim_[0] * float(pixel_dist_list[0]),
                                                                    dim_[1] * float(pixel_dist_list[1]),
                                                                    dim_[2] * float(pixel_dist_list[2])))
    log_f_.close()
    if not make_air:
        save_name_ = "{}_{}_{}.vox".format(model_name_, scale_, time_stamp_)
    else:
        save_name_ = "{}_air_{}_{}.vox".format(model_name_, scale_, time_stamp_)
    save_path_ = os.path.join(model_save_dir_path_, save_name_)

    print("*" * 60)
    print("vox file save path is {}".format(save_path_))
    vox_f_ = open("../tmp", "w")
    vox_f_.write(save_path_)
    vox_f_.close()

    f_ = open(save_path_, "w")
    f_.write('[SECTION VOXELS phantomER]\n')
    f_.write(str(nx_) + ' ' + str(ny_) + ' ' + str(nz_) + ' No. OF VOXELS IN X,Y,Z\n')
    f_.write(str(dx_) + ' ' + str(dy_) + ' ' + str(dz_) + ' VOXEL SIZE (cm) ALONG X,Y,Z\n')
    f_.write('1 COLUMN NUMBER WHERE MATERIAL ID IS LOCATED\n')
    f_.write('2 COLUMN NUMBER WHERE THE MASS DENSITY IS LOCATED\n')
    f_.write('0 BLANK LINES AT END OF X,Y-CYCLES (1=YES,0=NO)\n')
    f_.write('[END OF VXH SECTION]\n')

    print("start to generate .vox file")
    if not make_air:
        print("phantom_ has following values:{}".format(np.unique(phantom_)))
        for i in tqdm(range(total_voxels_)):
            # todo user use this to define the material vox
            if phantom_[i] == 0:  # 空气
                f_.write('1 ' + str(0.00120479) + '\n')
            elif phantom_[i] == 1:  # 脂肪
                f_.write('2 ' + str(0.95000000) + '\n')
            elif phantom_[i] == 2:  # barin
                f_.write('3 ' + str(1.05000000) + '\n')
            elif phantom_[i] == 3:  # 骨
                f_.write('4 ' + str(1.92000000) + '\n')
            elif phantom_[i] == 4:  # 肺部/脑脊液
                f_.write('5 ' + str(1.06000000) + '\n')
            # if phantom_[i] == 0:  # air
            #     f_.write('1 ' + str(0.00120479) + '\n')
            # elif phantom_[i] == 1:  # water
            #     f_.write('2 ' + str(1.0) + '\n')
            else:
                # f_.write('1 ' + str(0.00120479) + '\n')
                raise RuntimeError("there is a value {} in phantom,"
                                   "please check the code and notice "
                                   "if there are mistakes value between"
                                   "density and its no".format(phantom_[i]))
    else:
        for _ in tqdm(range(total_voxels_)):
            f_.write('1 ' + str(0.00120479) + '\n')
    f_.close()
    print("generation .vox file done")
    print("*" * 60)
    print("\n")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="define the model name")
    parser.add_argument("-name", "--model_name", default="model1109", type=str)
    parser.add_argument("-scale", "--scale", default=0.7, type=float)
    parser.add_argument("-pix_dist", "--pixel_dist", default="0.5, 0.5, 1", type=str)
    parser.add_argument("-air", "--make_air", default=False, type=bool)
    parser.add_argument("-nii", "--nii_file_path",
                        default=r"/mnt/520/dcp/demo/vox/medical_2/data/Head_Phantom.nii",
                        type=str)
    args = parser.parse_args()
    print(args)

    # time_stamp = add_timestamp()  # optional
    time_stamp = ""  # optional
    create_3d_vox_file(scale_=args.scale, model_name_=args.model_name,
                       pixel_dist=args.pixel_dist, time_stamp_=time_stamp,
                       nii_file_path_=args.nii_file_path,
                       make_air=args.make_air
                       )
    # pixel distance refers to [X, Y, Z]
